package com.LeeCode;

import java.util.HashMap;
import java.util.Map;

/**
 * 统计中位数为 K 的子数组
 */

public class Code2488 {
    public static void main(String[] args) {
        int[] nums = {2, 5, 1, 4, 3, 6};
        int k = 1;
        System.out.println(new Code2488().countSubarrays(nums, k));
    }

    public int countSubarrays(int[] nums, int k) {
        int kIndex = 0, n = nums.length;
        // 找到k的位置
        for (int i = 0; i < n; i++) {
            if (nums[i] == k) {
                kIndex = i;
                break;
            }
        }

        int preSum = 0;
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        // 计算k左边的前缀和
        for (int i = kIndex - 1; i >= 0; i--) {
            preSum += nums[i] > k ? 1 : -1;
            map.merge(preSum, 1, Integer::sum);
        }

        int ans = 0;
        preSum = 0;
        // 仅包含kIndex左边的部分
        ans += map.getOrDefault(0, 0); // 奇数长度
        ans += map.getOrDefault(1, 0); // 偶数长度
        // 包含k及右边元素的子数组
        for (int i = kIndex + 1; i < n; i++) {
            preSum += nums[i] > k ? 1 : -1;
            ans += map.getOrDefault(-preSum, 0); // 奇数长度
            ans += map.getOrDefault(-preSum + 1, 0); // 偶数长度
        }
        return ans;
    }
}
